% 质点-弹簧模型
% Gitee Repo

clc
clear

global n
global dx

L = 3;
dx = 0.5;
dt = 0.05;

[y0,x0] = meshgrid(-L:dx:L);
n = size(x0,1);
v0 = zeros(n,n,3);
r0 = zeros(n,n,3);
r0(:,:,1) = x0;
r0(:,:,2) = y0;
r0(:,:,3) = 0;

clear x0
clear y0

function fmat = compute_fmat_one_direction(x1,y1,z1,x2,y2,z2,dx)
    global n

    rvec(:,:,1) = x2-x1;
    rvec(:,:,2) = y2-y1;
    rvec(:,:,3) = z2-z1;

    dist = sqrt(sum(rvec.^2,3));
    rvec = rvec ./ (dist + 0.0001);

    for i = 1:3
        a = dist - dx;
        a(a>0) *= 10;
        a(a<0) *= 50;
        fmat(:,:,i) = a.*rvec(:,:,i);
    end

    return
end

function fmat = compute_fmat(r)
    global n
    global dx

    x = r(:,:,1);
    y = r(:,:,2);
    z = r(:,:,3);

    fmat = zeros(n,n,3);
    
    % x-y弹簧
    fmat(1:n-1,1:n,:) += compute_fmat_one_direction(x(1:n-1,1:n),y(1:n-1,1:n),z(1:n-1,1:n),x(2:n,1:n),y(2:n,1:n),z(2:n,1:n),dx);
    fmat(2:n,1:n,:) += compute_fmat_one_direction(x(2:n,1:n),y(2:n,1:n),z(2:n,1:n),x(1:n-1,1:n),y(1:n-1,1:n),z(1:n-1,1:n),dx);

    fmat(1:n,1:n-1,:) += compute_fmat_one_direction(x(1:n,1:n-1),y(1:n,1:n-1),z(1:n,1:n-1),x(1:n,2:n),y(1:n,2:n),z(1:n,2:n),dx);
    fmat(1:n,2:n,:) += compute_fmat_one_direction(x(1:n,2:n),y(1:n,2:n),z(1:n,2:n),x(1:n,1:n-1),y(1:n,1:n-1),z(1:n,1:n-1),dx);

    % 对角线弹簧
    fmat(1:n-1,1:n-1,:) += compute_fmat_one_direction(x(1:n-1,1:n-1),y(1:n-1,1:n-1),z(1:n-1,1:n-1),x(2:n,2:n),y(2:n,2:n),z(2:n,2:n),sqrt(2)*dx);
    fmat(2:n,2:n,:) += compute_fmat_one_direction(x(2:n,2:n),y(2:n,2:n),z(2:n,2:n),x(1:n-1,1:n-1),y(1:n-1,1:n-1),z(1:n-1,1:n-1),sqrt(2)*dx);

    fmat(2:n,1:n-1,:) += compute_fmat_one_direction(x(2:n,1:n-1),y(2:n,1:n-1),z(2:n,1:n-1),x(1:n-1,2:n),y(1:n-1,2:n),z(1:n-1,2:n),sqrt(2)*dx);
    fmat(1:n-1,2:n,:) += compute_fmat_one_direction(x(1:n-1,2:n),y(1:n-1,2:n),z(1:n-1,2:n),x(2:n,1:n-1),y(2:n,1:n-1),z(2:n,1:n-1),sqrt(2)*dx);
    return;
end


figure
clf

for tick = 0:10000
    fmat = compute_fmat(r0);
    fmat(n,floor(0.4*n):floor(0.6*n),1) -= 0.25*dx;
##    fmat(n,floor(0.4*n):floor(0.6*n),2) += 0.5*dx;
##    fmat(n,floor(0.4*n):floor(0.6*n),3) += -0.25*dx;
    fmat(n,floor(0.4*n):floor(0.6*n),3) += 0.05*(2*rand(n,n)-1)(n,floor(0.4*n):floor(0.6*n));
    v_half = v0 + fmat*dt/2;

    v_half(1,:,:) = 0;

    r1 = r0 + v_half*dt;

    fmat = compute_fmat(r1);
    v1 = v_half + fmat*dt/2;

    r0 = r1;
    v0 = 0.9*v1;
##    v0(:) = 0;

    if mod(tick,500) == 1
        clf
        hold on
        axis equal
        axis([-5,5,-5,5,-5,5])
        scatter3(r1(:,:,1),r1(:,:,2),r1(:,:,3));
        drawnow
        pause(0.1);
    end

    continue
end



